import psycopg2  # 导入PostgreSQL连接库
import pandas as pd  # 导入数据处理库
import time  # 导入时间库
from openai import OpenAI  # 导入OpenAI API
import re 

# 初始化OpenAI API客户端
client = OpenAI(
    #api_key="sk-FB9VFg4kAJfs4qCDLlQ10XRxJeNVzNIenLDJw8V9YqH1gzRF",
    api_key="sk-iPl7Bg48RtGNO2eVmBjDYhjm7rieL9vCLKAGIRawtBUExHCR",
    base_url="https://api.chatanywhere.tech/v1"
)

# 数据库连接参数
connection_params = {
    'dbname': 'imdb_database',
    'user': 'postgres',
    #'password': '0818',
    'password': 'postgres',
    'host': 'localhost',
    #'port': '5432'
    'port': '5434'
}

def connect_to_db(params):
    # 连接到数据库并返回连接对象
    try:
        connection = psycopg2.connect(**params)
        print("成功连接到数据库")
        return connection
    except Exception as e:
        print(f"连接数据库失败: {e}")
        return None

def generate_query_prompt():
    # 生成SQL查询的提示信息
#    prompt = """
#Generate 20 SQL queries based on the following table structure. 
#Ensure all queries are syntactically correct for PostgreSQL, with JOINs on at least five tables and filter conditions on columns.
#Do not include sequence numbers, comments, or extraneous characters like quotes or symbols. Only return the SQL queries.
#
#table structure：
#1. aka_name 表：
#   - id：整数，非空，主键
#   - person_id：整数，非空
#   - name：可变字符类型
#   - imdb_index：可变字符类型(3)
#   - name_pcode_cf：可变字符类型(11)
#   - name_pcode_nf：可变字符类型(11)
#   - surname_pcode：可变字符类型(11)
#   - md5sum：可变字符类型(65)
#
#2. aka_title 表：
#   - id：整数，非空，主键
#   - movie_id：整数，非空
#   - title：可变字符类型
#   - imdb_index：可变字符类型(4)
#   - kind_id：整数，非空
#   - production_year：整数
#   - phonetic_code：可变字符类型(5)
#   - episode_of_id：整数
#   - season_nr：整数
#   - episode_nr：整数
#   - note：可变字符类型(72)
#   - md5sum：可变字符类型(32)
#
#3. cast_info 表：
#   - id：整数，非空，主键
#   - person_id：整数，非空
#   - movie_id：整数，非空
#   - person_role_id：整数
#   - note：可变字符类型
#   - nr_order：整数
#   - role_id：整数，非空
#
#4. char_name 表：
#   - id：整数，非空，主键
#   - name：可变字符类型，非空
#   - imdb_index：可变字符类型(2)
#   - imdb_id：整数
#   - name_pcode_nf：可变字符类型(5)
#   - surname_pcode：可变字符类型(5)
#   - md5sum：可变字符类型(32)
#
# 5. comp_cast_type 表：
#    - id：整数，非空，主键
#    - kind：可变字符类型(32)，非空
#
#6. company_name 表：
#   - id：整数，非空，主键
#   - name：可变字符类型，非空
#   - country_code：可变字符类型(6)
#   - imdb_id：整数
#   - name_pcode_nf：可变字符类型(5)
#   - name_pcode_sf：可变字符类型(5)
#   - md5sum：可变字符类型(32)
#
#7. company_type 表：
#   - id：整数，非空，主键
#   - kind：可变字符类型(32)
#
#8. complete_cast 表：
#   - id：整数，非空，主键
#   - movie_id：整数
#   - subject_id：整数，非空
#   - status_id：整数，非空
#
#9. info_type 表：
#   - id：整数，非空，主键
#   - info：可变字符类型(32)，非空
#
#10. keyword 表：
#    - id：整数，非空，主键
#    - keyword：可变字符类型，非空
#    - phonetic_code：可变字符类型(5)
#
#11. kind_type 表：
#    - id：整数，非空，主键
#    - kind：可变字符类型(15)
#
#12. link_type 表：
#    - id：整数，非空，主键
#    - link：可变字符类型(32)，非空
#
#13. movie_companies 表：
#    - id：整数，非空，主键
#    - movie_id：整数，非空
#    - company_id：整数，非空
#    - company_type_id：整数，非空
#    - note：可变字符类型
#
#14. movie_info_idx 表：
#    - id：整数，非空，主键
#    - movie_id：整数，非空
#    - info_type_id：整数，非空
#    - info：可变字符类型，非空
#    - note：可变字符类型(1)
#
#15. movie_keyword 表：
#    - id：整数，非空，主键
#    - movie_id：整数，非空
#    - keyword_id：整数，非空
#
#16. movie_link 表：
#    - id：整数，非空，主键
#    - movie_id：整数，非空
#    - linked_movie_id：整数，非空
#    - link_type_id：整数，非空
#
#17. name 表：
#    - id：整数，非空，主键
#    - name：可变字符类型，非空
#    - imdb_index：可变字符类型(9)
#    - imdb_id：整数
#    - gender：可变字符类型(1)
#    - name_pcode_cf：可变字符类型(5)
#    - name_pcode_nf：可变字符类型(5)
#    - surname_pcode：可变字符类型(5)
#    - md5sum：可变字符类型(32)
#
#18. role_type 表：
#    - id：整数，非空，主键
#    - role：可变字符类型(32)，非空
#
#19. title 表：
#    - id：整数，非空，主键
#    - title：可变字符类型，非空
#    - imdb_index：可变字符类型(5)
#    - kind_id：整数，非空
#    - production_year：整数
#    - imdb_id：整数
#    - phonetic_code：可变字符类型(5)
#    - episode_of_id：整数
#    - season_nr：整数
#    - episode_nr：整数
#    - series_years：可变字符类型(49)
#    - md5sum：可变字符类型(32)
#
#20. movie_info 表：
#    - id：整数，非空，主键
#    - movie_id：整数，非空
#    - info_type_id：整数，非空
#    - info：可变字符类型，非空
#    - note：可变字符类型
#
#21. person_info 表：
#    - id：整数，非空，主键
#    - person_id：整数，非空
#    - info_type_id：整数，非空
#    - info：可变字符类型，非空
#    - note：可变字符类型
#
#"""

#2.
#Generate 20 SQL queries based on the following table structure. 
#Avoid using unnecessary characters, such as backticks or extraneous symbols. 
#Incorporate JOINs on at least five tables, different aggregation functions, and subqueries to increase complexity. 
#Do not include sequence numbers, comments, or any characters that could make the SQL invalid. Only return the SQL queries.
#Ensure all queries are syntactically correct for PostgreSQL, with JOINs on at least five tables and filter conditions on columns.

    prompt = """
    
Generate 20 SQL queries based on the following table structure. 
Ensure all queries are syntactically correct for PostgreSQL, with JOINs on at least five tables and filter conditions on columns.
Do not include sequence numbers, comments, or extraneous characters like quotes or symbols. Only return the SQL queries.

table structure：
1. aka_name Table：
   - id: Integer, not null, primary key
   - person_id: Integer, not null
   - name: Variable character type
   - imdb_index: Variable character type (3)
   - name_pcode_cf: Variable character type (11)
   - name_pcode_nf: Variable character type (11)
   - surname_pcode: Variable character type (11)
   - md5sum: Variable character type (65)

2. aka_title Table：
   - id: Integer, not null, primary key
   - movie_id: Integer, not null
   - title: Variable character type
   - imdb_index: Variable character type (4)
   - kind_id: Integer, not null
   - production_year: Integer
   - phonetic_code: Variable character type (5)
   - episode_of_id: Integer
   - season_nr: Integer
   - episode_nr: Integer
   - note: Variable character type (72)
   - md5sum: Variable character type (32)

3. cast_info Table：
   - id: Integer, not null, primary key
   - person_id: Integer, not null
   - movie_id: Integer, not null
   - person_role_id: Integer
   - note: Variable character type
   - nr_order: Integer
   - role_id: Integer, not null

4. char_name Table：
   - id: Integer, not null, primary key
   - name: Variable character type, not null
   - imdb_index: Variable character type (2)
   - imdb_id: Integer
   - name_pcode_nf: Variable character type (5)
   - surname_pcode: Variable character type (5)
   - md5sum: Variable character type (32)

 5. comp_cast_type Table：
    - id: Integer, not null, primary key
    - kind: Variable character type (32), not null

6. company_name Table：
   - id: Integer, not null, primary key
   - name: Variable character type, not null
   - country_code: Variable character type (6)
   - imdb_id: Integer
   - name_pcode_nf: Variable character type (5)
   - name_pcode_sf: Variable character type (5)
   - md5sum: Variable character type (32)

7. company_type Table：
   - id: Integer, not null, primary key
   - kind: Variable character type (32)

8. complete_cast Table：
   - id: Integer, not null, primary key
   - movie_id: Integer
   - subject_id: Integer, not null
   - status_id: Integer, not null

9. info_type Table：
   - id: Integer, not null, primary key
   - info: Variable character type (32), not null

10. keyword Table：
    - id: Integer, not null, primary key
    - keyword: Variable character type, not null
    - phonetic_code: Variable character type (5)

11. kind_type Table：
    - id: Integer, not null, primary key
    - kind: Variable character type (15)

12. link_type Table：
    - id: Integer, not null, primary key
    - link: Variable character type (32), not null

13. movie_companies Table：
    - id: Integer, not null, primary key
    - movie_id: Integer, not null
    - company_id: Integer, not null
    - company_type_id: Integer, not null
    - note: Variable character type

14. movie_info_idx Table：
    - id: Integer, not null, primary key
    - movie_id: Integer, not null
    - info_type_id: Integer, not null
    - info: Variable character type, not null
    - note: Variable character type (1)

15. movie_keyword Table：
    - id: Integer, not null, primary key
    - movie_id: Integer, not null
    - keyword_id: Integer, not null

16. movie_link Table：
    - id: Integer, not null, primary key
    - movie_id: Integer, not null
    - linked_movie_id: Integer, not null
    - link_type_id: Integer, not null

17. name Table：
    - id: Integer, not null, primary key
    - name: Variable character type, not null
    - imdb_index: Variable character type (9)
    - imdb_id: Integer
    - gender: Variable character type (1)
    - name_pcode_cf: Variable character type (5)
    - name_pcode_nf: Variable character type (5)
    - surname_pcode: Variable character type (5)
    - md5sum: Variable character type (32)

18. role_type Table：
    - id: Integer, not null, primary key
    - role: Variable character type (32), not null

19. title Table：
    - id: Integer, not null, primary key
    - title: Variable character type, not null
    - imdb_index: Variable character type (5)
    - kind_id: Integer, not null
    - production_year: Integer
    - imdb_id: Integer
    - phonetic_code: Variable character type (5)
    - episode_of_id: Integer
    - season_nr: Integer
    - episode_nr: Integer
    - series_years: Variable character type (49)
    - md5sum: Variable character type (32)

20. movie_info Table：
    - id: Integer, not null, primary key
    - movie_id: Integer, not null
    - info_type_id: Integer, not null
    - info: Variable character type, not null
    - note: Variable character type

21. person_info Table：
    - id: Integer, not null, primary key
    - person_id: Integer, not null
    - info_type_id: Integer, not null
    - info: Variable character type, not null
    - note: Variable character type

"""
    return prompt

def count_joins(query):
    # 计算查询中的JOIN数量和参与的表数
    join_count = 0
    tables_count = 0
    lines = query.split('\n')  # 按行分割查询
    for line in lines:
        clean_line = ' '.join(line.split())  # 清除多余空格
        if 'JOIN' in clean_line:
            join_count += 1  # 统计JOIN数量
            tables_count += 2  # 假设每个JOIN至少涉及两张表
    return join_count, tables_count

def clean_query(query):
    # 清理SQL查询，移除多余字符和注释
    cleaned_query = query.strip()  # 去除首尾空格
    cleaned_query = cleaned_query.split('--', 1)[0].strip()  # 移除注释部分
    cleaned_query = re.sub(r'^`?\s*\d+\.\s*', '', cleaned_query) #移除多余反引号和序号
    cleaned_query = ' '.join(cleaned_query.split())  # 移除多余空格
    return cleaned_query, count_joins(cleaned_query)  # 返回清理后的查询和JOIN统计

def generate_queries(prompt, num_queries=20):
    # 调用OpenAI API生成SQL查询
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=8192,
        n=1
    )
    
    raw_content = response.choices[0].message.content.strip()  # 获取API返回内容
    print(raw_content+'\n')
    queries = [clean_query(query) for query in raw_content.split(';') if query.strip()]  # 清理每个查询
    
    return queries[:num_queries]  # 返回指定数量的查询

def validate_query(cursor, connection, query):
    # 验证SQL查询的合法性
    try:
        cursor.execute(f"EXPLAIN {query}")  # 执行EXPLAIN查看查询计划
        explain_output = cursor.fetchall()  # 获取输出
        connection.commit()  # 提交事务
        return True, explain_output  # 返回有效性和输出
    except Exception as e:
        connection.rollback()  # 回滚事务
        return False, str(e)  # 返回无效性和错误信息

def correct_query_with_llm(failed_query, error_message):
    # 使用OpenAI API纠正错误的SQL查询
    prompt = f"Correct the following SQL query for PostgreSQL based on the error:\n\nQuery: {failed_query}\nError: {error_message}\n"
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=8192,
        n=1
    )
    
    corrected_query = response.choices[0].message.content.strip()  # 获取修正后的查询
    return clean_query(corrected_query)  # 返回清理后的修正查询

def main():
    # 主函数
    connection = connect_to_db(connection_params)  # 连接数据库
    if not connection:
        return  # 如果连接失败，结束程序

    cursor = connection.cursor()  # 创建游标
    
    prompt = generate_query_prompt()  # 生成查询提示
    
    start_time = time.time()  # 记录开始时间
    queries = generate_queries(prompt)  # 生成SQL查询
    end_time = time.time()  # 记录结束时间

    # 初始化统计变量
    total_joins = 0
    total_tables = 0
    for query, join_count in queries:
        total_joins += join_count[0]  # 统计总JOIN数量
        total_tables += join_count[1]  # 统计总表数

    # 计算平均生成时间、JOIN数量和表数
    avg_generation_time = (end_time - start_time) / len(queries)
    avg_joins = total_joins / len(queries) if queries else 0
    avg_tables = total_tables / len(queries) if queries else 0

    #log_file_path = r"D:\database\queries20.sql"  # 日志文件路径
    log_file_path = r"D:\IMDB\queries6.sql"  # 日志文件路径
    with open(log_file_path, 'w') as f:
        for idx, (query, join_count) in enumerate(queries):
            f.write(f"-- Query {idx + 1}\n")  # 写入查询序号
            f.write(f"-- Join Count: {join_count[0]}\n")  # 写入JOIN数量
            f.write(f"-- Tables Count: {join_count[1]}\n")  # 写入表数
            f.write(query + ";\n\n")  # 写入查询内容

    results = []  # 初始化结果列表
    for idx, (query, join_count) in enumerate(queries):
        valid, explain_output_or_error = validate_query(cursor, connection, query)  # 验证查询
        
        if not valid:
            # 如果无效，尝试纠正查询
            corrected_query, corrected_join_count = correct_query_with_llm(query, explain_output_or_error)
            valid, explain_output_or_error = validate_query(cursor, connection, corrected_query)
            if valid:
                validation_result = "Passed after correction"  # 修正后通过
                corrected_query = clean_query(corrected_query)  # 清理修正后的查询
                join_count = corrected_join_count  # 更新JOIN统计
            else:
                validation_result = "Failed after correction"  # 修正后仍然失败
        else:
            validation_result = "Passed"  # 原查询通过
        
        # 将结果添加到结果列表
        results.append({
            "Query ID": idx + 1,
            "Generated SQL Query": query,
            "Corrected SQL Query": corrected_query if not valid else "",
            "Join Count": join_count[0],
            "Tables Count": join_count[1],
            "Generation Time": avg_generation_time,
            "Validation Result": validation_result,
            "EXPLAIN Output": explain_output_or_error if validation_result.startswith("Passed") else "",
            "Error Message": explain_output_or_error if validation_result.startswith("Failed") else "",
            "Prompt": prompt,
        })

    df = pd.DataFrame(results)  # 创建结果数据框
    #excel_file_path = r"C:\Users\Aphro\Desktop\query_log20.xlsx"  # Excel文件路径
    excel_file_path = r"C:\Users\18375\Desktop\query_log6.xlsx"  # Excel文件路径
    df.to_excel(excel_file_path, index=False)  # 保存结果到Excel

    cursor.close()  # 关闭游标
    connection.close()  # 关闭数据库连接
    print(f"查询已保存到 {log_file_path}")  # 输出日志文件路径
    print(f"查询结果已保存到 {excel_file_path}")  # 输出Excel文件路径
    print(f"平均生成时间: {avg_generation_time} 秒")  # 输出平均生成时间
    print(f"平均 JOIN 子句数量: {avg_joins}")  # 输出平均JOIN数量
    print(f"平均参与 JOIN 操作的表数: {avg_tables}")  # 输出平均参与表数

if __name__ == "__main__":
    main()  # 运行主函数

